{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Interpretability with Monte Carlo Tree search\n",
    "\n",
    "Based on the paper Jin et al., [Multi-Objective Molecule Generation using Interpretable Substructures](https://arxiv.org/abs/2002.03244) and modified from Chemprop v1 [interpret.py](https://github.com/chemprop/chemprop/blob/master/chemprop/interpret.py)\n",
    "\n",
    "Please scroll to after the helper functions to change the model and data input and run the interpretation algorithm\n",
    "\n",
    "Note: \n",
    "- The interpret function does not yet work with additional atom or bond features, as the substructure extracted doesn't necessarily have the corresponding additional atom or bond features readily available.\n",
    "- It currently only works with single molecule model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/chemprop/chemprop/blob/main/examples/interpreting_monte_carlo_tree_search.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install chemprop from GitHub if running in Google Colab\n",
    "import os\n",
    "\n",
    "if os.getenv(\"COLAB_RELEASE_TAG\"):\n",
    "    try:\n",
    "        import chemprop\n",
    "    except ImportError:\n",
    "        !git clone https://github.com/chemprop/chemprop.git\n",
    "        %cd chemprop\n",
    "        !pip install .\n",
    "        %cd examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass, field\n",
    "import math\n",
    "from pathlib import Path\n",
    "import time\n",
    "from typing import Callable, Union, Iterable\n",
    "\n",
    "from lightning import pytorch as pl\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from rdkit import Chem\n",
    "import torch\n",
    "\n",
    "from chemprop import data, featurizers, models\n",
    "from chemprop.models import MPNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define helper function to make model predictions from SMILES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_prediction(\n",
    "    models: list[MPNN],\n",
    "    trainer: pl.Trainer,\n",
    "    smiles: list[str],\n",
    ") -> np.ndarray:\n",
    "    \"\"\"Makes predictions on a list of SMILES.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    models : list\n",
    "        A list of models to make predictions with.\n",
    "    smiles : list\n",
    "        A list of SMILES to make predictions on.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    list[list[float]]\n",
    "        A list of lists containing the predicted values.\n",
    "    \"\"\"\n",
    "\n",
    "    test_data = [data.MoleculeDatapoint.from_smi(smi) for smi in smiles]\n",
    "    test_dset = data.MoleculeDataset(test_data)\n",
    "    test_loader = data.build_dataloader(\n",
    "        test_dset, batch_size=1, num_workers=0, shuffle=False\n",
    "    )\n",
    "\n",
    "    with torch.inference_mode():\n",
    "        sum_preds = []\n",
    "        for model in models:\n",
    "            predss = trainer.predict(model, test_loader)\n",
    "            preds = torch.cat(predss, 0)\n",
    "            preds = preds.cpu().numpy()\n",
    "            sum_preds.append(preds)\n",
    "\n",
    "        # Ensemble predictions\n",
    "        sum_preds = sum(sum_preds)\n",
    "        avg_preds = sum_preds / len(models)\n",
    "\n",
    "    return avg_preds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Classes/functions relevant to Monte Carlo Tree Search\n",
    "\n",
    "Mostly similar to the scripts from Chemprop v1 [interpret.py](https://github.com/chemprop/chemprop/blob/master/chemprop/interpret.py) with additional documentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class MCTSNode:\n",
    "    \"\"\"Represents a node in a Monte Carlo Tree Search.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    smiles : str\n",
    "        The SMILES for the substructure at this node.\n",
    "    atoms : list\n",
    "        A list of atom indices in the substructure at this node.\n",
    "    W : float\n",
    "        The total action value, which indicates how likely the deletion will lead to a good rationale.\n",
    "    N : int\n",
    "        The visit count, which indicates how many times this node has been visited. It is used to balance exploration and exploitation.\n",
    "    P : float\n",
    "        The predicted property score of the new subgraphs' after the deletion, shown as R in the original paper.\n",
    "    \"\"\"\n",
    "\n",
    "    smiles: str\n",
    "    atoms: Iterable[int]\n",
    "    W: float = 0\n",
    "    N: int = 0\n",
    "    P: float = 0\n",
    "    children: list[...] = field(default_factory=list)\n",
    "\n",
    "    def __post_init__(self):\n",
    "        self.atoms = set(self.atoms)\n",
    "\n",
    "    def Q(self) -> float:\n",
    "        \"\"\"\n",
    "        Returns\n",
    "        -------\n",
    "        float\n",
    "            The mean action value of the node.\n",
    "        \"\"\"\n",
    "        return self.W / self.N if self.N > 0 else 0\n",
    "\n",
    "    def U(self, n: int, c_puct: float = 10.0) -> float:\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        n : int\n",
    "            The sum of the visit count of this node's siblings.\n",
    "        c_puct : float\n",
    "            A constant that controls the level of exploration.\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        float\n",
    "            The exploration value of the node.\n",
    "        \"\"\"\n",
    "        return c_puct * self.P * math.sqrt(n) / (1 + self.N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_clusters(mol: Chem.Mol) -> tuple[list[tuple[int, ...]], list[list[int]]]:\n",
    "    \"\"\"Finds clusters within the molecule. Jin et al. from [1]_ only allows deletion of one peripheral non-aromatic bond or one peripheral ring from each state,\n",
    "    so the clusters here are defined as non-ring bonds and the smallest set of smallest rings.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    mol : RDKit molecule\n",
    "        The molecule to find clusters in.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    tuple\n",
    "        A tuple containing:\n",
    "        - list of tuples: Each tuple contains atoms in a cluster.\n",
    "        - list of int: Each atom's cluster index.\n",
    "    \n",
    "    References\n",
    "    ----------\n",
    "    .. [1] Jin, Wengong, Regina Barzilay, and Tommi Jaakkola. \"Multi-objective molecule generation using interpretable substructures.\" International conference on machine learning. PMLR, 2020. https://arxiv.org/abs/2002.03244\n",
    "    \"\"\"\n",
    "\n",
    "    n_atoms = mol.GetNumAtoms()\n",
    "    if n_atoms == 1:  # special case\n",
    "        return [(0,)], [[0]]\n",
    "\n",
    "    clusters = []\n",
    "    for bond in mol.GetBonds():\n",
    "        a1 = bond.GetBeginAtom().GetIdx()\n",
    "        a2 = bond.GetEndAtom().GetIdx()\n",
    "        if not bond.IsInRing():\n",
    "            clusters.append((a1, a2))\n",
    "\n",
    "    ssr = [tuple(x) for x in Chem.GetSymmSSSR(mol)]\n",
    "    clusters.extend(ssr)\n",
    "\n",
    "    atom_cls = [[] for _ in range(n_atoms)]\n",
    "    for i in range(len(clusters)):\n",
    "        for atom in clusters[i]:\n",
    "            atom_cls[atom].append(i)\n",
    "\n",
    "    return clusters, atom_cls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_subgraph_from_mol(mol: Chem.Mol, selected_atoms: set[int]) -> tuple[Chem.Mol, list[int]]:\n",
    "    \"\"\"Extracts a subgraph from an RDKit molecule given a set of atom indices.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    mol : RDKit molecule\n",
    "        The molecule from which to extract a subgraph.\n",
    "    selected_atoms : list of int\n",
    "        The indices of atoms which form the subgraph to be extracted.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    tuple\n",
    "        A tuple containing:\n",
    "        - RDKit molecule: The subgraph.\n",
    "        - list of int: Root atom indices from the selected indices.\n",
    "    \"\"\"\n",
    "\n",
    "    selected_atoms = set(selected_atoms)\n",
    "    roots = []\n",
    "    for idx in selected_atoms:\n",
    "        atom = mol.GetAtomWithIdx(idx)\n",
    "        bad_neis = [y for y in atom.GetNeighbors() if y.GetIdx() not in selected_atoms]\n",
    "        if len(bad_neis) > 0:\n",
    "            roots.append(idx)\n",
    "\n",
    "    new_mol = Chem.RWMol(mol)\n",
    "\n",
    "    for atom_idx in roots:\n",
    "        atom = new_mol.GetAtomWithIdx(atom_idx)\n",
    "        atom.SetAtomMapNum(1)\n",
    "        aroma_bonds = [\n",
    "            bond for bond in atom.GetBonds() if bond.GetBondType() == Chem.rdchem.BondType.AROMATIC\n",
    "        ]\n",
    "        aroma_bonds = [\n",
    "            bond\n",
    "            for bond in aroma_bonds\n",
    "            if bond.GetBeginAtom().GetIdx() in selected_atoms\n",
    "            and bond.GetEndAtom().GetIdx() in selected_atoms\n",
    "        ]\n",
    "        if len(aroma_bonds) == 0:\n",
    "            atom.SetIsAromatic(False)\n",
    "\n",
    "    remove_atoms = [\n",
    "        atom.GetIdx() for atom in new_mol.GetAtoms() if atom.GetIdx() not in selected_atoms\n",
    "    ]\n",
    "    remove_atoms = sorted(remove_atoms, reverse=True)\n",
    "    for atom in remove_atoms:\n",
    "        new_mol.RemoveAtom(atom)\n",
    "\n",
    "    return new_mol.GetMol(), roots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_subgraph(smiles: str, selected_atoms: set[int]) -> tuple[str, list[int]]:\n",
    "    \"\"\"Extracts a subgraph from a SMILES given a set of atom indices.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    smiles : str\n",
    "        The SMILES string from which to extract a subgraph.\n",
    "    selected_atoms : list of int\n",
    "        The indices of atoms which form the subgraph to be extracted.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    tuple\n",
    "        A tuple containing:\n",
    "        - str: SMILES representing the subgraph.\n",
    "        - list of int: Root atom indices from the selected indices.\n",
    "    \"\"\"\n",
    "    # try with kekulization\n",
    "    mol = Chem.MolFromSmiles(smiles)\n",
    "    Chem.Kekulize(mol)\n",
    "    subgraph, roots = extract_subgraph_from_mol(mol, selected_atoms)\n",
    "    try:\n",
    "        subgraph = Chem.MolToSmiles(subgraph, kekuleSmiles=True)\n",
    "        subgraph = Chem.MolFromSmiles(subgraph)\n",
    "    except Exception:\n",
    "        subgraph = None\n",
    "\n",
    "    mol = Chem.MolFromSmiles(smiles)  # de-kekulize\n",
    "    if subgraph is not None and mol.HasSubstructMatch(subgraph):\n",
    "        return Chem.MolToSmiles(subgraph), roots\n",
    "\n",
    "    # If fails, try without kekulization\n",
    "    subgraph, roots = extract_subgraph_from_mol(mol, selected_atoms)\n",
    "    subgraph = Chem.MolToSmiles(subgraph)\n",
    "    subgraph = Chem.MolFromSmiles(subgraph)\n",
    "\n",
    "    if subgraph is not None:\n",
    "        return Chem.MolToSmiles(subgraph), roots\n",
    "    else:\n",
    "        return None, None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mcts_rollout(\n",
    "    node: MCTSNode,\n",
    "    state_map: dict[str, MCTSNode],\n",
    "    orig_smiles: str,\n",
    "    clusters: list[set[int]],\n",
    "    atom_cls: list[set[int]],\n",
    "    nei_cls: list[set[int]],\n",
    "    scoring_function: Callable[[list[str]], list[float]],\n",
    "    min_atoms: int = 15,\n",
    "    c_puct: float = 10.0,\n",
    ") -> float:\n",
    "    \"\"\"A Monte Carlo Tree Search rollout from a given MCTSNode.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    node : MCTSNode\n",
    "        The MCTSNode from which to begin the rollout.\n",
    "    state_map : dict\n",
    "        A mapping from SMILES to MCTSNode.\n",
    "    orig_smiles : str\n",
    "        The original SMILES of the molecule.\n",
    "    clusters : list\n",
    "        Clusters of atoms.\n",
    "    atom_cls : list\n",
    "        Atom indices in the clusters.\n",
    "    nei_cls : list\n",
    "        Neighboring cluster indices.\n",
    "    scoring_function : function\n",
    "        A function for scoring subgraph SMILES using a Chemprop model.\n",
    "    min_atoms : int\n",
    "        The minimum number of atoms in a subgraph.\n",
    "    c_puct : float\n",
    "        The constant controlling the level of exploration.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        The score of this MCTS rollout.\n",
    "    \"\"\"\n",
    "    # Return if the number of atoms is less than the minimum\n",
    "    cur_atoms = node.atoms\n",
    "    if len(cur_atoms) <= min_atoms:\n",
    "        return node.P\n",
    "\n",
    "    # Expand if this node has never been visited\n",
    "    if len(node.children) == 0:\n",
    "        # Cluster indices whose all atoms are present in current subgraph\n",
    "        cur_cls = set([i for i, x in enumerate(clusters) if x <= cur_atoms])\n",
    "\n",
    "        for i in cur_cls:\n",
    "            # Leaf atoms are atoms that are only involved in one cluster.\n",
    "            leaf_atoms = [a for a in clusters[i] if len(atom_cls[a] & cur_cls) == 1]\n",
    "\n",
    "            # This checks\n",
    "            # 1. If there is only one neighbor cluster in the current subgraph (so that we don't produce unconnected graphs), or\n",
    "            # 2. If the cluster has only two atoms and the current subgraph has only one leaf atom.\n",
    "            # If either of the conditions is met, remove the leaf atoms in the current cluster.\n",
    "            if len(nei_cls[i] & cur_cls) == 1 or len(clusters[i]) == 2 and len(leaf_atoms) == 1:\n",
    "                new_atoms = cur_atoms - set(leaf_atoms)\n",
    "                new_smiles, _ = extract_subgraph(orig_smiles, new_atoms)\n",
    "                if new_smiles in state_map:\n",
    "                    new_node = state_map[new_smiles]  # merge identical states\n",
    "                else:\n",
    "                    new_node = MCTSNode(new_smiles, new_atoms)\n",
    "                if new_smiles:\n",
    "                    node.children.append(new_node)\n",
    "\n",
    "        state_map[node.smiles] = node\n",
    "        if len(node.children) == 0:\n",
    "            return node.P  # cannot find leaves\n",
    "\n",
    "        scores = scoring_function([x.smiles for x in node.children])\n",
    "        for child, score in zip(node.children, scores):\n",
    "            child.P = score\n",
    "\n",
    "    sum_count = sum(c.N for c in node.children)\n",
    "    selected_node = max(node.children, key=lambda x: x.Q() + x.U(sum_count, c_puct=c_puct))\n",
    "    v = mcts_rollout(\n",
    "        selected_node,\n",
    "        state_map,\n",
    "        orig_smiles,\n",
    "        clusters,\n",
    "        atom_cls,\n",
    "        nei_cls,\n",
    "        scoring_function,\n",
    "        min_atoms=min_atoms,\n",
    "        c_puct=c_puct,\n",
    "    )\n",
    "    selected_node.W += v\n",
    "    selected_node.N += 1\n",
    "\n",
    "    return v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mcts(\n",
    "    smiles: str,\n",
    "    scoring_function: Callable[[list[str]], list[float]],\n",
    "    n_rollout: int,\n",
    "    max_atoms: int,\n",
    "    prop_delta: float,\n",
    "    min_atoms: int = 15,\n",
    "    c_puct: int = 10,\n",
    ") -> list[MCTSNode]:\n",
    "    \"\"\"Runs the Monte Carlo Tree Search algorithm.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    smiles : str\n",
    "        The SMILES of the molecule to perform the search on.\n",
    "    scoring_function : function\n",
    "        A function for scoring subgraph SMILES using a Chemprop model.\n",
    "    n_rollout : int\n",
    "        The number of MCTS rollouts to perform.\n",
    "    max_atoms : int\n",
    "        The maximum number of atoms allowed in an extracted rationale.\n",
    "    prop_delta : float\n",
    "        The minimum required property value for a satisfactory rationale.\n",
    "    min_atoms : int\n",
    "        The minimum number of atoms in a subgraph.\n",
    "    c_puct : float\n",
    "        The constant controlling the level of exploration.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    list\n",
    "        A list of rationales each represented by a MCTSNode.\n",
    "    \"\"\"\n",
    "\n",
    "    mol = Chem.MolFromSmiles(smiles)\n",
    "\n",
    "    clusters, atom_cls = find_clusters(mol)\n",
    "    nei_cls = [0] * len(clusters)\n",
    "    for i, cls in enumerate(clusters):\n",
    "        nei_cls[i] = [nei for atom in cls for nei in atom_cls[atom]]\n",
    "        nei_cls[i] = set(nei_cls[i]) - {i}\n",
    "        clusters[i] = set(list(cls))\n",
    "    for a in range(len(atom_cls)):\n",
    "        atom_cls[a] = set(atom_cls[a])\n",
    "\n",
    "    root = MCTSNode(smiles, set(range(mol.GetNumAtoms())))\n",
    "    state_map = {smiles: root}\n",
    "    for _ in range(n_rollout):\n",
    "        mcts_rollout(\n",
    "            root,\n",
    "            state_map,\n",
    "            smiles,\n",
    "            clusters,\n",
    "            atom_cls,\n",
    "            nei_cls,\n",
    "            scoring_function,\n",
    "            min_atoms=min_atoms,\n",
    "            c_puct=c_puct,\n",
    "        )\n",
    "\n",
    "    rationales = [\n",
    "        node\n",
    "        for _, node in state_map.items()\n",
    "        if len(node.atoms) <= max_atoms and node.P >= prop_delta\n",
    "    ]\n",
    "\n",
    "    return rationales"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "chemprop_dir = Path.cwd().parent\n",
    "model_path = (\n",
    "    chemprop_dir / \"tests\" / \"data\" / \"example_model_v2_regression_mol.pt\"\n",
    ")  # path to model checkpoint (.ckpt) or model file (.pt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MPNN(\n",
       "  (message_passing): BondMessagePassing(\n",
       "    (W_i): Linear(in_features=86, out_features=300, bias=False)\n",
       "    (W_h): Linear(in_features=300, out_features=300, bias=False)\n",
       "    (W_o): Linear(in_features=372, out_features=300, bias=True)\n",
       "    (dropout): Dropout(p=0.0, inplace=False)\n",
       "    (tau): ReLU()\n",
       "    (V_d_transform): Identity()\n",
       "    (graph_transform): GraphTransform(\n",
       "      (V_transform): Identity()\n",
       "      (E_transform): Identity()\n",
       "    )\n",
       "  )\n",
       "  (agg): MeanAggregation()\n",
       "  (bn): BatchNorm1d(300, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (predictor): RegressionFFN(\n",
       "    (ffn): MLP(\n",
       "      (0): Sequential(\n",
       "        (0): Linear(in_features=300, out_features=300, bias=True)\n",
       "      )\n",
       "      (1): Sequential(\n",
       "        (0): ReLU()\n",
       "        (1): Dropout(p=0.0, inplace=False)\n",
       "        (2): Linear(in_features=300, out_features=1, bias=True)\n",
       "      )\n",
       "    )\n",
       "    (criterion): MSE(task_weights=[[1.0]])\n",
       "    (output_transform): UnscaleTransform()\n",
       "  )\n",
       "  (X_d_transform): Identity()\n",
       "  (metrics): ModuleList(\n",
       "    (0-1): 2 x MSE(task_weights=[[1.0]])\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mpnn = models.MPNN.load_from_file(model_path)  # this is a dummy model for testing purposes\n",
    "mpnn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load data to run interpretation for"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "chemprop_dir = Path.cwd().parent\n",
    "test_path = chemprop_dir / \"tests\" / \"data\" / \"regression\" / \"mol\" / \"mol.csv\"\n",
    "smiles_column = \"smiles\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>smiles</th>\n",
       "      <th>lipo</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Cn1c(CN2CCN(CC2)c3ccc(Cl)cc3)nc4ccccc14</td>\n",
       "      <td>3.54</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>COc1cc(OC)c(cc1NC(=O)CSCC(=O)O)S(=O)(=O)N2C(C)...</td>\n",
       "      <td>-1.18</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>COC(=O)[C@@H](N1CCc2sccc2C1)c3ccccc3Cl</td>\n",
       "      <td>3.69</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>OC[C@H](O)CN1C(=O)C(Cc2ccccc12)NC(=O)c3cc4cc(C...</td>\n",
       "      <td>3.37</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Cc1cccc(C[C@H](NC(=O)c2cc(nn2C)C(C)(C)C)C(=O)N...</td>\n",
       "      <td>3.10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>95</th>\n",
       "      <td>CC(C)N(CCCNC(=O)Nc1ccc(cc1)C(C)(C)C)C[C@H]2O[C...</td>\n",
       "      <td>2.20</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>96</th>\n",
       "      <td>CCN(CC)CCCCNc1ncc2CN(C(=O)N(Cc3cccc(NC(=O)C=C)...</td>\n",
       "      <td>2.04</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>97</th>\n",
       "      <td>CCSc1c(Cc2ccccc2C(F)(F)F)sc3N(CC(C)C)C(=O)N(C)...</td>\n",
       "      <td>4.49</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>98</th>\n",
       "      <td>COc1ccc(Cc2c(N)n[nH]c2N)cc1</td>\n",
       "      <td>0.20</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>99</th>\n",
       "      <td>CCN(CCN(C)C)S(=O)(=O)c1ccc(cc1)c2cnc(N)c(n2)C(...</td>\n",
       "      <td>2.00</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>100 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                               smiles  lipo\n",
       "0             Cn1c(CN2CCN(CC2)c3ccc(Cl)cc3)nc4ccccc14  3.54\n",
       "1   COc1cc(OC)c(cc1NC(=O)CSCC(=O)O)S(=O)(=O)N2C(C)... -1.18\n",
       "2              COC(=O)[C@@H](N1CCc2sccc2C1)c3ccccc3Cl  3.69\n",
       "3   OC[C@H](O)CN1C(=O)C(Cc2ccccc12)NC(=O)c3cc4cc(C...  3.37\n",
       "4   Cc1cccc(C[C@H](NC(=O)c2cc(nn2C)C(C)(C)C)C(=O)N...  3.10\n",
       "..                                                ...   ...\n",
       "95  CC(C)N(CCCNC(=O)Nc1ccc(cc1)C(C)(C)C)C[C@H]2O[C...  2.20\n",
       "96  CCN(CC)CCCCNc1ncc2CN(C(=O)N(Cc3cccc(NC(=O)C=C)...  2.04\n",
       "97  CCSc1c(Cc2ccccc2C(F)(F)F)sc3N(CC(C)C)C(=O)N(C)...  4.49\n",
       "98                        COc1ccc(Cc2c(N)n[nH]c2N)cc1  0.20\n",
       "99  CCN(CCN(C)C)S(=O)(=O)c1ccc(cc1)c2cnc(N)c(n2)C(...  2.00\n",
       "\n",
       "[100 rows x 2 columns]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_test = pd.read_csv(test_path)\n",
    "df_test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "HPU available: False, using: 0 HPUs\n",
      "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/setup.py:177: GPU available but not used. You can set it by doing `Trainer(accelerator='gpu')`.\n"
     ]
    }
   ],
   "source": [
    "trainer = pl.Trainer(logger=None, enable_progress_bar=False, accelerator=\"cpu\", devices=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Running interpretation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MCTS options\n",
    "rollout = 10  # number of MCTS rollouts to perform. If mol.GetNumAtoms() > 50, consider setting n_rollout = 1 to avoid long computation time\n",
    "\n",
    "c_puct = 10.0  # constant that controls the level of exploration\n",
    "\n",
    "max_atoms = 20  # maximum number of atoms allowed in an extracted rationale\n",
    "\n",
    "min_atoms = 8  # minimum number of atoms in an extracted rationale\n",
    "\n",
    "prop_delta = 0.5  # Minimum score to count as positive.\n",
    "# In this algorithm, if the predicted property from the substructure if larger than prop_delta, the substructure is considered satisfactory.\n",
    "# This value depends on the property you want to interpret. 0.5 is a dummy value for demonstration purposes\n",
    "\n",
    "num_rationales_to_keep = 5  # number of rationales to keep for each molecule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the scoring function. \"Score\" for a substructure is the predicted property value of the substructure.\n",
    "\n",
    "models = [mpnn]\n",
    "\n",
    "property_for_interpretation = \"lipo\"\n",
    "\n",
    "property_id = (\n",
    "    df_test.columns.get_loc(property_for_interpretation) - 1\n",
    ")  # property index in the dataset; -1 for the SMILES column\n",
    "\n",
    "\n",
    "def scoring_function(smiles: list[str]) -> list[float]:\n",
    "    return make_prediction(\n",
    "        models=models,\n",
    "        trainer=trainer,\n",
    "        smiles=smiles,\n",
    "    )[:, property_id]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Cn1c(CN2CCN(CC2)c3ccc(Cl)cc3)nc4ccccc14',\n",
       " 'COc1cc(OC)c(cc1NC(=O)CSCC(=O)O)S(=O)(=O)N2C(C)CCc3ccccc23',\n",
       " 'COC(=O)[C@@H](N1CCc2sccc2C1)c3ccccc3Cl',\n",
       " 'OC[C@H](O)CN1C(=O)C(Cc2ccccc12)NC(=O)c3cc4cc(Cl)sc4[nH]3',\n",
       " 'Cc1cccc(C[C@H](NC(=O)c2cc(nn2C)C(C)(C)C)C(=O)NCC#N)c1']"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# only use the first 5 SMILES for demonstration purposes\n",
    "all_smiles = df_test[smiles_column].tolist()[:5]\n",
    "all_smiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:424: The 'predict_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=7` in the `DataLoader` to improve performance.\n",
      "[19:04:06] Can't kekulize mol.  Unkekulized atoms: 10 11 12 13 14\n",
      "[19:04:06] Can't kekulize mol.  Unkekulized atoms: 11 12 13 14 15\n",
      "[19:04:06] Can't kekulize mol.  Unkekulized atoms: 8 9 10 11 12\n",
      "[19:04:06] Can't kekulize mol.  Unkekulized atoms: 7 8 9 10 11\n",
      "[19:04:06] Can't kekulize mol.  Unkekulized atoms: 1 2 3 4 5\n",
      "[19:04:06] Can't kekulize mol.  Unkekulized atoms: 0 1 3 4 5\n",
      "[19:04:06] Can't kekulize mol.  Unkekulized atoms: 0 1 2 3 4\n",
      "[19:04:06] Can't kekulize mol.  Unkekulized atoms: 11 12 13 14 15\n",
      "[19:04:06] Can't kekulize mol.  Unkekulized atoms: 8 9 10 11 12\n",
      "[19:04:06] Can't kekulize mol.  Unkekulized atoms: 7 8 9 10 11\n",
      "[19:04:06] Can't kekulize mol.  Unkekulized atoms: 10 11 12 13 14\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 13 s, sys: 1.38 s, total: 14.4 s\n",
      "Wall time: 3.67 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "results_df = {\"smiles\": [], property_for_interpretation: []}\n",
    "\n",
    "for i in range(num_rationales_to_keep):\n",
    "    results_df[f\"rationale_{i}\"] = []\n",
    "    results_df[f\"rationale_{i}_score\"] = []\n",
    "\n",
    "for smiles in all_smiles:\n",
    "    score = scoring_function([smiles])[0]\n",
    "    if score > prop_delta:\n",
    "        rationales = mcts(\n",
    "            smiles=smiles,\n",
    "            scoring_function=scoring_function,\n",
    "            n_rollout=rollout,\n",
    "            max_atoms=max_atoms,\n",
    "            prop_delta=prop_delta,\n",
    "            min_atoms=min_atoms,\n",
    "            c_puct=c_puct,\n",
    "        )\n",
    "    else:\n",
    "        rationales = []\n",
    "\n",
    "    results_df[\"smiles\"].append(smiles)\n",
    "    results_df[property_for_interpretation].append(score)\n",
    "\n",
    "    if len(rationales) == 0:\n",
    "        for i in range(num_rationales_to_keep):\n",
    "            results_df[f\"rationale_{i}\"].append(None)\n",
    "            results_df[f\"rationale_{i}_score\"].append(None)\n",
    "    else:\n",
    "        min_size = min(len(x.atoms) for x in rationales)\n",
    "        min_rationales = [x for x in rationales if len(x.atoms) == min_size]\n",
    "        rats = sorted(min_rationales, key=lambda x: x.P, reverse=True)\n",
    "\n",
    "        for i in range(num_rationales_to_keep):\n",
    "            if i < len(rats):\n",
    "                results_df[f\"rationale_{i}\"].append(rats[i].smiles)\n",
    "                results_df[f\"rationale_{i}_score\"].append(rats[i].P)\n",
    "            else:\n",
    "                results_df[f\"rationale_{i}\"].append(None)\n",
    "                results_df[f\"rationale_{i}_score\"].append(None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>smiles</th>\n",
       "      <th>lipo</th>\n",
       "      <th>rationale_0</th>\n",
       "      <th>rationale_0_score</th>\n",
       "      <th>rationale_1</th>\n",
       "      <th>rationale_1_score</th>\n",
       "      <th>rationale_2</th>\n",
       "      <th>rationale_2_score</th>\n",
       "      <th>rationale_3</th>\n",
       "      <th>rationale_3_score</th>\n",
       "      <th>rationale_4</th>\n",
       "      <th>rationale_4_score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Cn1c(CN2CCN(CC2)c3ccc(Cl)cc3)nc4ccccc14</td>\n",
       "      <td>2.253542</td>\n",
       "      <td>c1ccc2c(c1)n[cH:1][nH:1]2</td>\n",
       "      <td>2.275024</td>\n",
       "      <td>None</td>\n",
       "      <td>NaN</td>\n",
       "      <td>None</td>\n",
       "      <td>NaN</td>\n",
       "      <td>None</td>\n",
       "      <td>NaN</td>\n",
       "      <td>None</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>COc1cc(OC)c(cc1NC(=O)CSCC(=O)O)S(=O)(=O)N2C(C)...</td>\n",
       "      <td>2.235016</td>\n",
       "      <td>O=[SH:1]c1c[cH:1][cH:1]cc1[OH:1]</td>\n",
       "      <td>2.252582</td>\n",
       "      <td>c1c([OH:1])c([S:1][NH2:1])c[cH:1][cH:1]1</td>\n",
       "      <td>2.252185</td>\n",
       "      <td>c1c(N[CH3:1])[cH:1]c[cH:1]c1[SH:1]</td>\n",
       "      <td>2.251068</td>\n",
       "      <td>c1c([S:1][NH2:1])[cH:1]cc([OH:1])[cH:1]1</td>\n",
       "      <td>2.250288</td>\n",
       "      <td>c1c([NH2:1])[cH:1]c[cH:1]c1[S:1][NH2:1]</td>\n",
       "      <td>2.249267</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>COC(=O)[C@@H](N1CCc2sccc2C1)c3ccccc3Cl</td>\n",
       "      <td>2.245891</td>\n",
       "      <td>c1cc[cH:1]c([CH2:1][CH2:1][OH:1])c1</td>\n",
       "      <td>2.249289</td>\n",
       "      <td>O=[CH:1][CH2:1]c1cccc[cH:1]1</td>\n",
       "      <td>2.249207</td>\n",
       "      <td>c1cc[cH:1]c([C@@H]([CH3:1])[NH2:1])c1</td>\n",
       "      <td>2.247827</td>\n",
       "      <td>Clc1ccccc1[CH2:1][NH2:1]</td>\n",
       "      <td>2.245391</td>\n",
       "      <td>Clc1ccccc1[CH2:1][CH3:1]</td>\n",
       "      <td>2.243280</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>OC[C@H](O)CN1C(=O)C(Cc2ccccc12)NC(=O)c3cc4cc(C...</td>\n",
       "      <td>2.249847</td>\n",
       "      <td>c1c([CH3:1])[nH]c2s[cH:1]cc12</td>\n",
       "      <td>2.267990</td>\n",
       "      <td>Clc1cc2c[cH:1][nH]c2s1</td>\n",
       "      <td>2.267004</td>\n",
       "      <td>O=C1N(C[CH3:1])[CH:1]=[CH:1]C[CH2:1]1</td>\n",
       "      <td>2.211323</td>\n",
       "      <td>None</td>\n",
       "      <td>NaN</td>\n",
       "      <td>None</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Cc1cccc(C[C@H](NC(=O)c2cc(nn2C)C(C)(C)C)C(=O)N...</td>\n",
       "      <td>2.228097</td>\n",
       "      <td>c1cc(C[CH2:1][NH2:1])c[cH:1]c1</td>\n",
       "      <td>2.247070</td>\n",
       "      <td>c1cc(C[CH2:1][CH3:1])c[cH:1]c1</td>\n",
       "      <td>2.245314</td>\n",
       "      <td>Cn1nc([CH3:1])cc1[CH2:1][NH2:1]</td>\n",
       "      <td>2.225729</td>\n",
       "      <td>C[CH2:1]c1cc([CH2:1][NH2:1])[nH:1]n1</td>\n",
       "      <td>2.223793</td>\n",
       "      <td>c1c([CH3:1])n[nH:1]c1[CH2:1]N[CH3:1]</td>\n",
       "      <td>2.223478</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                              smiles      lipo  \\\n",
       "0            Cn1c(CN2CCN(CC2)c3ccc(Cl)cc3)nc4ccccc14  2.253542   \n",
       "1  COc1cc(OC)c(cc1NC(=O)CSCC(=O)O)S(=O)(=O)N2C(C)...  2.235016   \n",
       "2             COC(=O)[C@@H](N1CCc2sccc2C1)c3ccccc3Cl  2.245891   \n",
       "3  OC[C@H](O)CN1C(=O)C(Cc2ccccc12)NC(=O)c3cc4cc(C...  2.249847   \n",
       "4  Cc1cccc(C[C@H](NC(=O)c2cc(nn2C)C(C)(C)C)C(=O)N...  2.228097   \n",
       "\n",
       "                           rationale_0  rationale_0_score  \\\n",
       "0            c1ccc2c(c1)n[cH:1][nH:1]2           2.275024   \n",
       "1     O=[SH:1]c1c[cH:1][cH:1]cc1[OH:1]           2.252582   \n",
       "2  c1cc[cH:1]c([CH2:1][CH2:1][OH:1])c1           2.249289   \n",
       "3        c1c([CH3:1])[nH]c2s[cH:1]cc12           2.267990   \n",
       "4       c1cc(C[CH2:1][NH2:1])c[cH:1]c1           2.247070   \n",
       "\n",
       "                                rationale_1  rationale_1_score  \\\n",
       "0                                      None                NaN   \n",
       "1  c1c([OH:1])c([S:1][NH2:1])c[cH:1][cH:1]1           2.252185   \n",
       "2              O=[CH:1][CH2:1]c1cccc[cH:1]1           2.249207   \n",
       "3                    Clc1cc2c[cH:1][nH]c2s1           2.267004   \n",
       "4            c1cc(C[CH2:1][CH3:1])c[cH:1]c1           2.245314   \n",
       "\n",
       "                             rationale_2  rationale_2_score  \\\n",
       "0                                   None                NaN   \n",
       "1     c1c(N[CH3:1])[cH:1]c[cH:1]c1[SH:1]           2.251068   \n",
       "2  c1cc[cH:1]c([C@@H]([CH3:1])[NH2:1])c1           2.247827   \n",
       "3  O=C1N(C[CH3:1])[CH:1]=[CH:1]C[CH2:1]1           2.211323   \n",
       "4        Cn1nc([CH3:1])cc1[CH2:1][NH2:1]           2.225729   \n",
       "\n",
       "                                rationale_3  rationale_3_score  \\\n",
       "0                                      None                NaN   \n",
       "1  c1c([S:1][NH2:1])[cH:1]cc([OH:1])[cH:1]1           2.250288   \n",
       "2                  Clc1ccccc1[CH2:1][NH2:1]           2.245391   \n",
       "3                                      None                NaN   \n",
       "4      C[CH2:1]c1cc([CH2:1][NH2:1])[nH:1]n1           2.223793   \n",
       "\n",
       "                               rationale_4  rationale_4_score  \n",
       "0                                     None                NaN  \n",
       "1  c1c([NH2:1])[cH:1]c[cH:1]c1[S:1][NH2:1]           2.249267  \n",
       "2                 Clc1ccccc1[CH2:1][CH3:1]           2.243280  \n",
       "3                                     None                NaN  \n",
       "4     c1c([CH3:1])n[nH:1]c1[CH2:1]N[CH3:1]           2.223478  "
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results_df = pd.DataFrame(results_df)\n",
    "results_df"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chemprop",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
